

#pragma once
#include <stdint.h>
#include <vector>
#include <thread>
#include <memory>
#include <assert.h>
#include <mutex>

#include "ClassBase.h"

namespace FlameEngine
{
	/////////////////////////////////////////////////
	class FrameAllocMemChunk
	{
	public:
		FrameAllocMemChunk(size_t size);
		~FrameAllocMemChunk();
		
		//allocate memory from chunk, nullptr if chunk is full
		void* Allocate(size_t size, size_t align, int32_t& allocByteCount);

		//reset memory chunk
		inline void Clear()
		{
			memset(m_pDataStart, 0, m_chunkSize);
			m_freeByteCount = m_chunkSize;
		}

		inline size_t GetChunkSize() { return m_chunkSize; }

	private:
		const size_t m_chunkSize;
		size_t m_freeByteCount;
		uint8_t* m_pDataStart;
		uint8_t* m_pDataCurrent;

		FrameAllocMemChunk(FrameAllocMemChunk&) = delete;
		FrameAllocMemChunk& operator=(FrameAllocMemChunk&) = delete;
	};

	/////////////////////////////////////////////////
	class FrameAllocator
	{
	public:
		FrameAllocator();

		~FrameAllocator();

		//create object from allocator
		template<typename T,typename... Args>
		T* CreateObject(int align, Args... params)
		{
			if(align == 0)
				align = 4;

			void* pNewMem = Allocate(sizeof(T), align);
			return (T*)(new(pNewMem)T(params...));
		}

		//create raw memory from allcator
		void* Allocate(size_t size, int32_t align = 4);

		//clear all data 
		void Clear();

		size_t GetAllocSize() { return m_chunks.size() * m_chunkSize; }

		size_t GetUsedSize() { return m_usedByteNum; }

	private:
		std::vector<FrameAllocMemChunk*> m_chunks;
		size_t m_usedByteNum;
		int32_t m_curChunkIdx;
		int32_t m_chunkSize;

		FrameAllocator(FrameAllocator&) = delete;
		FrameAllocator& operator= (FrameAllocator&) = delete;
	};

	/////////////////////////////////////////////////
	//FrameAllocatorThreadSafe guarantee allocate memory during 
	//a frame is thread safe
	class FrameAllocatorThreadSafe
	{
	public:
		FrameAllocatorThreadSafe();

		~FrameAllocatorThreadSafe();

		//create raw memory from allocator
		inline void* Allocate(size_t size, int32_t align = 4)
		{
			GetAllocator()->Allocate(size, align);
		}

		//create object from allocator
		template<typename T, typename... Args>
		inline T* CreateObject(int align, Args... params)
		{
			return GetAllocator()->CreateObject<T>(align,params...);
		}

		//clear all data,this function is not thread safe 
		//should only be called at the end of frame
		inline void Clear()
		{
			for(auto& alloc : m_threadAlloc)
				alloc.second->Clear();
		}

		size_t GetAllocSize()
		{
			size_t allocSize = 0;
			for(auto& alloc : m_threadAlloc)
				allocSize += alloc.second->GetAllocSize();
			
			return allocSize;
		}

		size_t GetUsedSize()
		{
			size_t usedSize = 0;
			for(auto& alloc : m_threadAlloc)
				usedSize += alloc.second->GetUsedSize();
			
			return usedSize;
		}

	private:
		//we won't have more then 10 threads,so vector is fine
		std::vector<std::pair<std::thread::id, FrameAllocator*>> m_threadAlloc;

		std::mutex m_mutex;

		FrameAllocator* GetAllocator();
	};

	/////////////////////////////////////////////////
	template<typename T>
	class TripleFrameAllocatorTemplate:public NonCopyAssignClass 
	{
	public:
		TripleFrameAllocatorTemplate() {}
		~TripleFrameAllocatorTemplate() {}

		T* GetCurrentAlloc()
		{
			return &m_frameAllocator[m_alloc0];
		}

		void SetCurrentFrame(uint32_t frameCount)
		{
			if(frameCount > m_frameCount)
			{
				//swap allocator
				m_frameAllocator[m_alloc2].Clear();
				uint32_t temp = m_alloc2;
				m_alloc2 = m_alloc1;
				m_alloc1 = m_alloc0;
				m_alloc0 = temp;
				m_frameCount = frameCount
			}
		}	
	private:
		//Triple buffering frame allocator
		T m_frameAllocator[3];
		uint32_t m_frameCount;
		uint32_t m_alloc0;
		uint32_t m_alloc1;
		uint32_t m_alloc2;
	};

	typedef TripleFrameAllocatorTemplate<FrameAllocator> TripleFrameAllocator;
	typedef TripleFrameAllocatorTemplate<FrameAllocatorThreadSafe> TripleFrameAllocatorThreadSafe;
}